# -*- coding:utf-8  -*-
# @Author: ChenYangMing
# @Time: 2024-08-29

import numpy as np
from sklearn.model_selection import train_test_split


class DataUtils:
    """
    数据工具类，提供数据处理的方法。
    包括删除行、删除列、添加行、添加列、删除重复项、填充缺失值、分割数据集和创建时间序列数据等操作。
    """
    @staticmethod
    def remove_rows(
            df,
            condition,
    ):
        """
        根据给定的条件删除DataFrame中的行。
        :param df: pd.DataFrame, 待处理的DataFrame。
        :param condition: function, 用于判断每行是否符合删除条件的函数。
        :return: pd.DataFrame, 删除指定行后的DataFrame。
        """
        return df[~df.apply(condition, axis=1)]

    @staticmethod
    def remove_columns(
            df,
            columns_to_remove,
    ):
        """
        从DataFrame中删除指定的
        :param df: pd.DataFrame, 待处理的DataFrame。
        :param columns_to_remove: list, 需要删除的列名列表。
        :return: pd.DataFrame, 删除指定列后的DataFrame。
        """
        return df.drop(columns=columns_to_remove, axis=1)

    @staticmethod
    def add_row(
            df,
            new_row,
            index,
    ):
        """
        向DataFrame中添加新的一行。

        :param df: pd.DataFrame, 待处理的DataFrame。
        :param new_row: pd.Series, 要添加的新行，索引应该与df的列对齐。
        :param index: object, 新行的索引值，如果为None，则添加到df的末尾。
        :return: pd.DataFrame, 添加新行后的DataFrame。
        """
        new_df = df.append(new_row, ignore_index=True)
        if index is not None:
            new_df = new_df.set_index(index)
        return new_df

    @staticmethod
    def add_column(
            df,
            new_column,
    ):
        """
        向DataFrame中添加新的一列。
        :param df: pd.DataFrame, 待处理的DataFrame。
        :param new_column: dict, 要添加的新列，键为列名，值为列的数据。
        :return: pd.DataFrame, 添加新列后的DataFrame。
        """
        return df.assign(**new_column)

    @staticmethod
    def remove_duplicates(
            df,
    ):
        """
        删除DataFrame中的重复行。
        :param df: pd.DataFrame, 待处理的DataFrame。
        :return: pd.DataFrame, 删除重复行后的DataFrame。
        """
        return df.drop_duplicates()

    @staticmethod
    def fill_missing(
            df,
            strategy='ffill',
            **kwargs,
    ):
        """
        填充DataFrame中的缺失值。
        :param df: pd.DataFrame, 待处理的DataFrame。
        :param strategy: str, 填充策略，'ffill'表示前向填充，'bfill'表示后向填充。
        :param kwargs: dict, 传递给DataFrame.fillna的其他参数。
        :return: pd.DataFrame, 填充缺失值后的DataFrame。
        """
        return df.fillna(method=strategy, **kwargs)

    @staticmethod
    def split_data(
            x,
            y,
            test_size=0.2,
            random_state=None,
    ):
        """
        将数据集划分为训练集和测试集。
        :param x: pd.DataFrame, 特征数据。
        :param y: pd.Series, 目标数据。
        :param test_size: float, 测试集的比例。
        :param random_state: int, 随机种子，用于确保结果的可重复性。
        :return: tuple, 包含训练集特征、测试集特征、训练集目标和测试集目标的元组。
        """
        x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=test_size, random_state=random_state)
        return x_train, x_test, y_train, y_test

    @staticmethod
    def create_time_series(
            x,
            hl,
            pl,
    ):
        """
        将数据转换为时间序列数据。
        :param x: list 或 np.ndarray, 输入特征数据。
        :param hl: int, 历史窗口长度。
        :param pl: int, 预测长度。
        :return: tuple, 包含特征集X和目标集y的元组，均为np.ndarray类型。
        """
        x, y = [], []
        total_time_window = hl + pl
        loop_count = len(x) - total_time_window + 1  # 循环次数
        for i in range(loop_count):
            x.append(x[i:(i + hl)])  # 历史数据
            y.append(x[i + hl:i + total_time_window])  # 目标
        return np.array(x), np.array(y)


# 调用示例，假设 df 是已经定义好的 DataFrame
# df = pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]})
# df = DataUtils.remove_columns(df, ['B'])
# df = DataUtils.fill_missing(df)
